from neo4j import GraphDatabase
from shared_enum import Layer
from frame_structure import Frame
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from frame_manager import FrameManager


class NeoManager:
    def __init__(self, URI, auth):
        self.layer = 'session'
        self.URI = URI
        self.auth = auth
        self.driver = None
        self.layer = Layer.Annotation

    def change_layer(self, layer):
        if isinstance(layer, Layer):
            self.layer = layer
        elif isinstance(layer, str):
            if layer == 'annotation':
                self.layer = Layer.Annotation
            elif layer == 'primary':
                self.layer = Layer.Primary
            elif layer == 'prototype':
                self.layer = Layer.Prototype
            else:
                raise ValueError
        return self

    def search(self, term: str, frame_manager) -> list[Frame]:
        frames = []
        with self.driver.session(database=self.layer.name) as session:
            res = session.run(
                '''
                MATCH (f:FrameInstance{frame:$name})
                RETURN f
                ''', name=term
            )
            _res = res.value()
            for node in _res:
                frames.append(Frame(frame_manager, term).create_from_db(self, node))
        return frames

    def __enter__(self):
        self.driver = GraphDatabase.driver(self.URI, auth=self.auth, keep_alive=True)
        return self

    def __exit__(self, exc_type, exc_value, trace):
        self.driver.close()